package com.woohua.thread.limite;

import java.util.Vector;

/***
 * @title CurrentLimiter
 * @description：限流器，可以添加限流规则，根据限流规则获取流量通行证，释放流量通行证；如果获取通行证失败则抛出异常
 * @author woo hua
 * @version 1.0.0
 * @create 2023/3/14 8:55
 **/
public class CurrentLimiter {

    /**
     * 本地线程变量，存储一次请求获取到的通行证，和其他并发请求隔离开，
     * 在controller执行完后释放本次请求获得的通行证
     */
    private static ThreadLocal<Vector<LimitRule>> localAcquiredLimitRules = new ThreadLocal<>();

    /**
     * 所有限流规则
     */
    private static Vector<LimitRule> allLimitRules = new Vector<>();

    /**
     * 私有构造器，避免实例化
     */
    private CurrentLimiter() {
    }

    /**
     * 添加限流规则，在Spring启动时添加，不需要加锁，如果在运行中添加，需要加锁
     *
     * @param rule
     */
    public static void addRule(LimitRule rule) {
        print(rule);
        allLimitRules.add(rule);
    }

    /**
     * 获取流量通行证，所有流量规则都要获取后才能通过，如果一个不能获取则抛出异常
     *
     * @param urlPath
     * @throws Exception
     */
    public static void tryAcquire(String urlPath) throws Exception {
        // 有限流规则则处理
        if (allLimitRules.size() > 0) {
            // 能获取到通行证的流量规则要保存下来，在controller执行完后要释放
            Vector<LimitRule> acquireLimitRules = new Vector<>();
            for (LimitRule rule : allLimitRules) {
                // 获取通行证
                AcquireResult acquireResult = rule.tryAcquire(urlPath);

                if (AcquireResult.ACQUIRE_SUCCESS == acquireResult.getResult()) {
                    acquireLimitRules.add(rule);
                    // 获取到通行证的流量规格添加到本地线程变量
                    localAcquiredLimitRules.set(acquireLimitRules);
                } else if (AcquireResult.ACQUIRE_FAILED == acquireResult.getResult()) {
                    // 如果获取不到通行证则抛出异常
                    StringBuffer buffer = new StringBuffer();
                    buffer.append("The request [").append(urlPath).append("] exceeds max num traffic limit, this limit is")
                            .append(rule.getMaxConcurrent()).append(", available permit is").append(acquireResult.getAvailablePermits())
                            .append(".");

                    System.out.println(buffer);
                    throw new Exception(buffer.toString());
                } else {
                    StringBuffer buffer = new StringBuffer();
                    buffer.append("This path does not match the limit rule, path is [").append(urlPath)
                            .append("], pattern is [").append(rule.getPattern()).append("].");
                    System.out.println(buffer);
                }
            }
        }
    }

    /**
     * 释放获取到的通行证。在controller执行完后调用（抛出异常也需要调用）
     */
    public static void release() {
        Vector<LimitRule> acquireLimitRules = localAcquiredLimitRules.get();
        if (null != acquireLimitRules && acquireLimitRules.size() > 0) {
            acquireLimitRules.forEach(rule -> {
                rule.release();
            });
        }

        // destroy 本地线程变量，避免内存泄露
        localAcquiredLimitRules.remove();
    }

    /**
     * 打印限流规则信息
     *
     * @param rule
     */
    private static void print(LimitRule rule) {
        StringBuffer buffer = new StringBuffer();
        buffer.append("Add Limit Rule, Max Concurrent").append(rule.getMaxConcurrent())
                .append(", Pattern: ").append(rule.getPattern());
        System.out.println(buffer.toString());
    }
}
